import os

import numpy as np
import torch
import torch.nn.functional as F

import utils
from agent.Nets import PolicyNet, ValueNet, PolicyNetContinuous
from agent.AgentBase import AgentBase


class PPO(AgentBase):
    """PPO-clip算法"""

    def __init__(self, state_dim, hidden_dim, action_dim, hidden_layers, actor_lr, critic_lr, lmbda, epochs, eps, gamma,
                 device, args,
                 action_type="discrete"):
        super().__init__(args)
        if action_type == "continuous":
            self.actor = PolicyNetContinuous(state_dim, hidden_dim, action_dim).to(device)
        else:
            self.actor = PolicyNet(state_dim, hidden_dim, action_dim, hidden_layers).to(device)
        self.critic = ValueNet(state_dim, hidden_dim, hidden_layers).to(device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=critic_lr)
        self.gamma = gamma
        self.lmbda = lmbda
        # 一条序列的数据用来训练轮数
        self.epochs = epochs
        # PPO-clip中截断范围的参数
        self.eps = eps
        self.device = device
        self.action_type = action_type

    def take_action(self, state, prediction=False):
        state = torch.tensor(np.array([state]), dtype=torch.float).to(self.device)
        if self.action_type == "continuous":
            mu, std = self.actor(state)
            action_list = torch.distributions.Normal(mu, std)
            return [action_list.sample().item()]
        else:
            probs = self.actor(state)
            action_list = torch.distributions.Categorical(probs)
            action = action_list.sample()
            return action.item()

    def update(self, transition_dict):
        states = torch.tensor(transition_dict['states'], dtype=torch.float).to(self.device)
        if self.action_type == "continuous":
            actions = torch.tensor(transition_dict['actions'], dtype=torch.float).view(-1, 1).to(self.device)
        else:
            actions = torch.tensor(transition_dict['actions'], dtype=torch.int64).view(-1, 1).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'], dtype=torch.float).view(-1, 1).to(self.device)
        next_states = torch.tensor(transition_dict['next_states'], dtype=torch.float).to(self.device)
        dones = torch.tensor(transition_dict['dones'], dtype=torch.float).view(-1, 1).to(self.device)

        td_target = rewards + self.gamma * self.critic(next_states) * (1 - dones)
        td_error = td_target - self.critic(states)
        # 计算优势
        advantage = utils.compute_advantage(self.gamma, self.lmbda, td_error.cpu()).to(self.device)
        if self.action_type == "continuous":
            mu, std = self.actor(states)
            action_dists = torch.distributions.Normal(mu.detach(), std.detach())
            # 动作是正态分布
            old_log_probs = action_dists.log_prob(actions)
        else:
            old_log_probs = torch.log(self.actor(states).gather(1, actions)).detach()

        for _ in range(self.epochs):
            if self.action_type == "continuous":
                mu, std = self.actor(states)
                action_dists = torch.distributions.Normal(mu, std)
                # 动作是正态分布
                log_probs = action_dists.log_prob(actions)
            else:
                log_probs = torch.log(self.actor(states).gather(1, actions))

            # import sampling
            ratio = torch.exp(log_probs - old_log_probs)
            surr1 = ratio * advantage
            surr2 = torch.clamp(ratio, 1 - self.eps, 1 + self.eps) * advantage  # 截断
            # PPO损失函数
            actor_loss = torch.mean(-torch.min(surr1, surr2))
            # 价值网络损失函数
            critic_loss = torch.mean(F.mse_loss(self.critic(states), td_target.detach()))
            self.actor_optimizer.zero_grad()
            self.critic_optimizer.zero_grad()
            actor_loss.backward()
            critic_loss.backward()
            self.actor_optimizer.step()
            self.critic_optimizer.step()

    def save(self):
        if not os.path.exists(self.args.model_path_pre):
            os.mkdir(self.args.model_path_pre)
        torch.save(self.actor.state_dict(), self.args.model_path)

    def load(self, args):
        torch.load(self.actor, self.args.model_path)


if __name__ == '__main__':
    from run.train import Trainer
    import gym
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--env_name", type=str, default="CartPole-v0")
    parser.add_argument("--agent_name", type=str)
    parser.add_argument("--actor_lr", type=float, default=1e-4)
    parser.add_argument("--critic_lr", type=float, default=5e-3)
    parser.add_argument("--lmbda", type=float, default=0.9)
    parser.add_argument("--internal_log", type=int, default=50, help="每隔多少个episode打印日志")
    parser.add_argument("--evaluate_internal", type=int, default=10, help="每隔多少个episode评估一下模型")
    parser.add_argument("--record_path_pre", type=str, default="../run/logs/", help="存储tensorboard日志文件的地方")
    parser.add_argument("--agent_type", type=str, default="onPolicy", help="Agent类型")
    parser.add_argument("--epochs", type=int, default=10, help="PPO中一条轨迹被用来更新的次数")
    parser.add_argument("--model_path_pre", type=str, default="../run/models/", help="模型保存位置以及名字")
    parser.add_argument("--sampling_strategy", type=str, default="", help="")
    parser.add_argument("--epsilon", type=float, default=0.1)
    parser.add_argument("--temperature", type=float, default=5, help="玻尔兹曼分布温度系数")
    parser.add_argument("--tau", type=float, default=0.005, help="软更新参数")
    parser.add_argument("--num_episodes", type=int, default=2000)
    args = parser.parse_args()

    env_names = ['CartPole-v1', 'LunarLander-v2']
    for env_name in env_names:
        args.agent_name = "PPO"
        args.model_path = args.model_path_pre + args.agent_name + "_" + env_name + "_" + args.sampling_strategy + ".pt"
        args.record_path = args.record_path_pre + args.agent_name + "_" + env_name + "_" + args.sampling_strategy
        args.env_name = env_name

        env = gym.make(env_name)
        state_dim = env.observation_space.shape[0]
        action_dim = env.action_space.n
        hidden_dim = 128
        hidden_layers = 1
        gamma = 0.99
        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        args.device = device

        agent = PPO(state_dim, hidden_dim, action_dim, args.actor_lr, args.critic_lr, args.lmbda, args.epochs,
                    args.epsilon, gamma, device, args)

        trainer = Trainer(env, agent, None, args)
        trainer.train_on_policy()
